#include "acl/acl.h"
#include "aclnn/acl_meta.h"
#include <vector>
#include <iostream>

#define CHECK_RET(cond, return_expr) \
    if (!(cond))                     \
    {                                \
        return_expr;                 \
    }
#define LOG_PRINT(message, ...) printf(message, ##__VA_ARGS__);
#define CHECK_SUCCESS_RETURN(ret, message) CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT(message "ERROR:%d\n", ret); return ret;)

#define CAST_TO_ACL_TENSOR(value) reinterpret_cast<NDTensor *>(value)->tensorPtr
#define CAST_TO_ENV(value) reinterpret_cast<aclEnv *>(value)
#define GET_STREAM(value) CAST_TO_ENV(value)->stream
#define ALLOC_WORKSPACE(workspaceSize)                                                                        \
    void *workspaceAddr = nullptr;                                                                            \
    if (workspaceSize > 0)                                                                                    \
    {                                                                                                         \
        ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);                          \
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret;); \
    }
class aclEnv
{
public:
    aclrtContext context = nullptr;
    aclrtStream stream = nullptr;
    int32_t device_id;

public:
    aclEnv(int32_t deviceId);
    virtual ~aclEnv();
    int syncStream();
};

class NDTensor
{
public:
    uint64_t size = 0;
    void *hostAddr = nullptr;
    void *deviceAddr = nullptr;
    aclTensor *tensorPtr = nullptr;

public:
    NDTensor(void *shape, uint64_t shape_size, int dataType, void *tensorData);
    virtual ~NDTensor();
    void to_cpu();
};
